{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from hloc import (\n",
    "    extract_features,\n",
    "    match_features,\n",
    "    reconstruction,\n",
    "    visualization,\n",
    "    pairs_from_retrieval,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "In this notebook, we will run SfM reconstruction from scratch on a set of images. We choose the [South-Building dataset](https://openaccess.thecvf.com/content_cvpr_2013/html/Hane_Joint_3D_Scene_2013_CVPR_paper.html) - we will download it later. First, we define some paths."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = Path(\"datasets/South-Building/images/\")\n",
    "\n",
    "outputs = Path(\"outputs/sfm/\")\n",
    "sfm_pairs = outputs / \"pairs-netvlad.txt\"\n",
    "sfm_dir = outputs / \"sfm_superpoint+superglue\"\n",
    "\n",
    "retrieval_conf = extract_features.confs[\"netvlad\"]\n",
    "feature_conf = extract_features.confs[\"superpoint_aachen\"]\n",
    "matcher_conf = match_features.confs[\"superglue\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the dataset\n",
    "The dataset is simply a set of images. The intrinsic parameters will be extracted from the EXIF data and refined with SfM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not images.exists():\n",
    "    !wget https://cvg-data.inf.ethz.ch/local-feature-evaluation-schoenberger2017/South-Building.zip -P datasets/\n",
    "    !unzip -q datasets/South-Building.zip -d datasets/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Find image pairs via image retrieval\n",
    "We extract global descriptors with NetVLAD and find for each image the most similar ones. For smaller dataset we can instead use exhaustive matching via `hloc/pairs_from_exhaustive.py`, which would find $\\frac{n(n-1)}{2}$ images pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieval_path = extract_features.main(retrieval_conf, images, outputs)\n",
    "pairs_from_retrieval.main(retrieval_path, sfm_pairs, num_matched=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract and match local features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_path = extract_features.main(feature_conf, images, outputs)\n",
    "match_path = match_features.main(\n",
    "    matcher_conf, sfm_pairs, feature_conf[\"output\"], outputs\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3D reconstruction\n",
    "Run COLMAP on the features and matches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = reconstruction.main(sfm_dir, images, sfm_pairs, feature_path, match_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization\n",
    "We visualize some of the registered images, and color their keypoint by visibility, track length, or triangulated depth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualization.visualize_sfm_2d(model, images, color_by=\"visibility\", n=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualization.visualize_sfm_2d(model, images, color_by=\"track_length\", n=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualization.visualize_sfm_2d(model, images, color_by=\"depth\", n=5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
